from selenium import webdriver
from selenium.webdriver.chrome.options import Options
from selenium.webdriver.common.by import By
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
import time
import logging
import os
import json
import requests
from datetime import datetime
from dotenv import load_dotenv

# 加载环境变量
load_dotenv()

# 配置常量
BALANCE_THRESHOLD = float(os.getenv('BALANCE_THRESHOLD', '100.0'))
WEBHOOK_URL = os.getenv('WEBHOOK_URL', '')

def setup_logger():
    """配置日志"""
    if not os.path.exists('logs'):
        os.makedirs('logs')
    
    log_filename = f"logs/balance_check_{datetime.now().strftime('%Y%m%d_%H%M%S')}.log"
    
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s [%(levelname)s] %(message)s',
        handlers=[
            logging.FileHandler(log_filename, encoding='utf-8'),
            logging.StreamHandler()
        ]
    )
    return logging.getLogger(__name__)

def format_balance(balance_text):
    """格式化余额显示"""
    try:
        amount = float(''.join(filter(lambda x: x.isdigit() or x == '.', balance_text)))
        formatted = f"￥{amount:,.2f}"
        return formatted, amount
    except:
        return balance_text, 0.0

def send_webhook_notification(balance_info, logger):
    """发送 webhook 通知"""
    if not WEBHOOK_URL:
        logger.warning("未配置 WEBHOOK_URL，跳过通知")
        return
        
    try:
        payload = {
            "timestamp": datetime.now().isoformat(),
            "balance_info": balance_info,
            "message": f"""BigModel 账户余额警告：
- 当前可用余额: {balance_info['available_balance']}
- 累计充值: {balance_info['total_recharge']}
- 赠送金额: {balance_info['gift_amount']}
- 预警阈值: ￥{balance_info['threshold']:,.2f}"""
        }
        
        response = requests.post(
            WEBHOOK_URL,
            json=payload,
            headers={'Content-Type': 'application/json'}
        )
        
        if response.status_code == 200:
            logger.info("Webhook 通知发送成功")
        else:
            logger.error(f"Webhook 通知发送失败: {response.status_code} - {response.text}")
            
    except Exception as e:
        logger.error(f"Webhook 通知发送异常: {str(e)}")

def check_bigmodel_balance(logger):
    """检查 BigModel 余额"""
    # 检查环境变量
    username = os.getenv('BIGMODEL_USERNAME')
    password = os.getenv('BIGMODEL_PASSWORD')
    
    if not username or not password:
        logger.error("未找到必要的环境变量！")
        logger.error("请确保 .env 文件中包含以下配置：")
        logger.error("BIGMODEL_USERNAME=你的用户名")
        logger.error("BIGMODEL_PASSWORD=你的密码")
        return None, 0.0
    
    if username == 'your_username' or password == 'your_password':
        logger.error("请修改 .env 文件中的默认用户名和密码！")
        return None, 0.0
    
    # 设置 Chrome 选项
    chrome_options = Options()
    chrome_options.add_argument('--no-sandbox')
    chrome_options.add_argument('--start-maximized')
    chrome_options.add_argument('lang=zh_CN')
    
    # 配置远程 WebDriver
    try:
        driver = webdriver.Remote(
            command_executor='http://localhost:4444',
            options=chrome_options
        )
    except Exception as e:
        logger.error(f"无法连接到 Selenium Grid: {str(e)}")
        logger.error("请确保 Selenium Grid 容器正在运行：")
        logger.error("docker run -d -p 4444:4444 -p 7900:7900 --shm-size='2g' selenium/standalone-chrome:latest")
        return None, 0.0
    
    try:
        # 访问登录页面
        logger.info("访问登录页面")
        driver.get('https://bigmodel.cn/login')
        
        # 等待登录表单加载
        wait = WebDriverWait(driver, 20)
        
        # 切换到账号登录模式
        logger.info("切换到账号登录模式")
        account_tab = wait.until(EC.element_to_be_clickable((By.ID, "tab-password")))
        account_tab.click()
        
        # 输入账号密码
        logger.info("输入账号密码")
        username_input = wait.until(EC.presence_of_element_located(
            (By.XPATH, "//input[@placeholder='请输入用户名/邮箱/手机号']")))
        password_input = driver.find_element(
            By.XPATH, "//input[@placeholder='请输入密码']")
        
        username_input.send_keys(username)
        password_input.send_keys(password)
        
        # 点击登录
        logger.info("点击登录按钮")
        login_button = driver.find_element(
            By.XPATH, "//button[contains(@class, 'loginDialog-btn')]")
        login_button.click()
        
        # 等待登录成功
        time.sleep(3)
        
        # 访问财务页面
        logger.info("访问财务页面")
        driver.get('https://bigmodel.cn/finance/overview')
        time.sleep(2)
        
        # 获取余额信息
        logger.info("获取余额信息")
        balance_area = wait.until(EC.presence_of_element_located(
            (By.XPATH, "//div[contains(@class, 'amount-area')]")))
        
        # 获取各项金额
        available_balance = balance_area.find_element(
            By.XPATH, ".//div[contains(text(), '可用余额')]/following-sibling::div[contains(@class, 'price')]").text
        total_recharge = balance_area.find_element(
            By.XPATH, ".//div[contains(text(), '累计充值')]/following-sibling::div[contains(@class, 'price')]").text
        gift_amount = balance_area.find_element(
            By.XPATH, ".//div[contains(text(), '赠送金额')]/following-sibling::div[contains(@class, 'price')]").text
        
        # 格式化金额
        formatted_balance, amount = format_balance(available_balance)
        formatted_recharge, _ = format_balance(total_recharge)
        formatted_gift, _ = format_balance(gift_amount)
        
        # 记录信息
        logger.info(f"""
账户余额信息:
- 可用余额: {formatted_balance}
- 累计充值: {formatted_recharge}
- 赠送金额: {formatted_gift}
        """.strip())
        
        # 检查余额阈值
        if amount < BALANCE_THRESHOLD:
            logger.warning(f"余额 {formatted_balance} 低于预警阈值 ￥{BALANCE_THRESHOLD:,.2f}")
            send_webhook_notification({
                'available_balance': formatted_balance,
                'total_recharge': formatted_recharge,
                'gift_amount': formatted_gift,
                'threshold': BALANCE_THRESHOLD
            }, logger)
        
        return formatted_balance, amount
        
    except Exception as e:
        logger.error(f"发生错误: {str(e)}")
        # 保存错误信息
        timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
        screenshot_file = f'logs/error_screenshot_{timestamp}.png'
        driver.save_screenshot(screenshot_file)
        logger.error(f"页面截图已保存至: {screenshot_file}")
        return None, 0.0
        
    finally:
        logger.info("关闭浏览器")
        driver.quit()

def main():
    """主函数"""
    logger = setup_logger()
    logger.info("开始执行余额查询")
    
    try:
        balance, amount = check_bigmodel_balance(logger)
        if balance:
            logger.info(f"余额查询完成: {balance}")
        return balance, amount
    except Exception as e:
        logger.error(f"程序执行失败: {str(e)}")
        return None, 0.0

if __name__ == "__main__":
    main() 